import pdb, code, decimal
import traceback
import pyodbc, sys, re, csv, os, datetime

sys.path.append('.')
pyodbc.pooling = False

class mssql_error(Exception):
    def __init__(self, message):
        Exception.__init__(self, message)

def build(server, admin_login, admin_password, rm_dbname, setup_path):
    h = handle(server=server, login=admin_login, password=admin_password, autocommit=True)
    h.build(rm_dbname = rm_dbname, setup_path = setup_path)

class handle():
    def __init__(self, server, login = None, password = None, database = None, autocommit=False, isolation = None):
        self.server = self.location = str(server)
        connection_params = ['DRIVER={SQL Server}', 'SERVER=' + self.server]
        self.login = login
        if login is not None and len(login) > 0:
            self.login = str(login)
            connection_params.append('UID=' + self.login)
            self.password = password
            if password is not None and len(password) > 0:
                self.password = str(password)
                connection_params.append('PWD=' + self.password)
        else:
            self.login = '\\'.join([os.environ.get('USERDOMAIN'), os.environ.get('USERNAME')])
            self.password = None

        self.database = database
        if database is not None and len(database) > 0:
            self.database = str(database)
            self.location = self.location + '\\' + self.database
            connection_params.append('DATABASE=' + self.database)
        
        connection_str = ';'.join(connection_params)
        self.connection = pyodbc.connect(connection_str, autocommit=autocommit)
        self.cursor = self.connection.cursor()
        if isolation is not None:
            self.cursor.execute("SET TRANSACTION ISOLATION LEVEL " + isolation)
            self.connection.commit()

    def __del__(self):
        try:
            self.connection.rollback()
            self.cursor.close()
            self.connection.close()
        except:
            pass

    def commit(self):
        self.connection.commit()
    
    def rollback(self):
        self.connection.rollback()

    def close(self):
        self.cursor.close()
        self.connection.close()

    def setAutocommit(self, bit):
        self.connection.autocommit = bit

    def execute(self, query):
        self.cursor.execute(query)

    def pendingTransactions(self):
        self.execute('SELECT @@TRANCOUNT')
        trancount = self.fetch()[0][0]
        return trancount > 0
    
    def fetch(self):
        return self.cursor.fetchall()

    def dbExists(self, db_name):
        self.execute("SELECT physical_name FROM sys.master_files WHERE database_id = DB_ID(N'" + db_name + "')")
        return len(self.fetch()) > 0

    def getTransactions(self, sql_path):
        with open(sql_path, 'r') as f:
            query_list = f.readlines()
        GO_inds = [ind for ind, line in enumerate(query_list) if re.match('^[ ]*GO$', line)]
        transactions = []
        last_GO = -1
        for ind in GO_inds:
            query = query_list[(last_GO+1):ind]
            query = ''.join(query)
            transactions.append(query)
            last_GO = ind
        if GO_inds[-1] != (len(query_list)-1):
            query = query_list[(GO_inds[-1]+1):len(query_list)]
            query = ''.join(query)
            transactions.append(query)
        return(transactions)

    def executeScript(self, sql_path):
        transactions = self.getTransactions(sql_path)
        self.setAutocommit(True)
        for transaction in transactions:
            try:
                self.execute(transaction)
            except:
                type, value, tb = sys.exc_info()
                exc_formatted = ''.join(traceback.format_exception_only(type, value))
                value = '\n\n'.join([exc_formatted, 'Offending transaction:', transaction])
                raise mssql_error(value)

    def build(self, rm_dbname, setup_path):
        current_commit_setting = self.connection.autocommit
        self.setAutocommit(True)
        if self.dbExists(rm_dbname):
            self.tearDownDB(rm_dbname)
        try:
            self.executeScript(setup_path)
        except:
            if self.dbExists('SWMMR'):
                self.tearDownDB('SWMMR')
            self.setAutocommit(current_commit_setting)
            raise
        self.setAutocommit(current_commit_setting)

    def diffCSV(self, table_name, csv_path):
        expected_rows = self.csvImportForExistingTable(csv_path, table_name)
        return self.diffTable(table_name=table_name, ordered_by='OBJECTID', rows=expected_rows)

    def tearDownDB(self, name):
        commit_setting = self.connection.autocommit # get current autocommit setting
        self.connection.autocommit = True # set autocommit to True
        self.cursor.execute("USE [MASTER]")
        self.cursor.execute("ALTER DATABASE [" + name + "] SET SINGLE_USER WITH ROLLBACK IMMEDIATE") # breaks any lingering connections
        self.cursor.execute("DROP DATABASE [" + name + "]")
        self.connection.autocommit = commit_setting # reset autocommit setting to what it was
        
    def getColumnNames(self, table_name):
        '''
        returns the columns for "table_name"
        '''
        split_name = table_name.split('.')
        if len(split_name) > 1:
            table_name = split_name[1]
        if table_name[0] == '#':
            self.execute("SELECT name FROM tempdb.sys.columns WHERE object_id = OBJECT_ID('tempdb.." + table_name + "')")
            columns = self.fetch()
            columns = [str(col[0]) for col in columns]
        else:
            try:
                self.cursor.columns(table=table_name)
            except:
                print("Error running DB.getColumnNames: " + "Can't find table '" + table_name + "'.")
                raise
            else:
                name_fieldnum = 3
                rows = self.cursor.fetchall()
                columns = [str(row[name_fieldnum]) for row in rows]

        return columns

    def getProcedureParameters(self, procedure_name):
        '''
        returns list of parameter names for procedure. Names include '@' prefix.
        "procedure" can be the procedure name with or without the schema prefix.
        '''
        split_name = procedure_name.split('.')
        if len(split_name) > 1:
            procedure_name = split_name[1]

        self.cursor.execute("SELECT PARAMETER_NAME FROM INFORMATION_SCHEMA.PARAMETERS WHERE SPECIFIC_NAME=" +
                quote(procedure_name) + " AND PARAMETER_MODE = 'IN'")
        parameter_names = self.cursor.fetchall()
        parameter_names = [str(name[0]) for name in parameter_names]
        return(parameter_names)

    def getLogins(self, default_db=None):
        query = "SELECT name FROM sys.server_principals WHERE type_desc = 'SQL_LOGIN' "
        if default_db is not None:
            query += 'AND default_database_name = ' + quote(default_db)
        self.cursor.execute(query)
        principals = self.cursor.fetchall()
        principals = [str(principal[0]) for principal in principals]
        return principals

    def getUsers(self):
        self.execute("SELECT name FROM sys.database_principals WHERE type_desc = 'SQL_USER'")
        principals = self.fetch()
        principals = [str(principal[0]) for principal in principals]
        return principals
    
    def getAllRows(self, table_name, order_by=None):
        '''
        Returns the rows in "table_name" as a list of dictionaries, where the keys in each row are 
        the column names of the table, eg., [{'col1' : 'A', col2 : 1}, {'col1' : 'B', 'col2' : 2}],
        with rows possibly ordered by the column provided in "order_by"

        If table is empty, returns an empty list
        '''
        query = 'SELECT * FROM ' + table_name
        colnames = self.getColumnNames(table_name)
        if order_by is not None and order_by in colnames:
            query = query + ' ORDER BY ' + order_by
        self.cursor.execute(query)
        all_rows = self.cursor.fetchall()
        all_rows = [list(row) for row in all_rows]
        all_rows = [dict(zip(colnames, row)) for row in all_rows]
        return(all_rows)

    def getFunctionReturnFieldNames(self, function_name):
        split_name = function_name.split('.')
        schema = split_name[0]
        function = split_name[1]
        self.cursor.execute(' '.join(["SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.ROUTINE_COLUMNS WHERE TABLE_SCHEMA =",
                                    quote(schema), "AND TABLE_NAME =", quote(function), "ORDER BY ORDINAL_POSITION"]))
        field_names = self.cursor.fetchall()
        field_names = [str(field_name[0]) for field_name in field_names]
        return field_names

    def functionCall(self, function, named_args=None):
        '''
        executes the user function "function" with arguments supplied in "named_args" 
        as a dictionary, and returns the return value of the user function.
        '''
        if named_args is None:
            call = function + '()'
        else:
            # combine arguments into appropriate call string
            self.checkParameterNames(function, named_args.keys())
            parameter_names = self.getProcedureParameters(function)
            args_ordered = [named_args[param_name] for param_name in parameter_names]
            # if an argument is a string, put it in quotes. If it's None, replace it with NULL
            for i in xrange(len(args_ordered)):
                arg = args_ordered[i]
                if isinstance(arg, unicode):
                    arg = str(arg)
                if isinstance(arg, str):
                    args_ordered[i] = quote(arg)
                elif arg is not None:
                    args_ordered[i] = str(arg)
                else:
                    args_ordered[i] = 'NULL'

            args = ','.join(args_ordered)
            call = function + '(' + args + ')'

        try:
            self.cursor.execute("SELECT " + call)
            result = self.cursor.fetchone()[0] # cursor returns row tuples
            return result
        except:
            # if function returns a table, previous select won't work
            try:
                self.cursor.execute("SELECT * FROM " + call)
                result = self.cursor.fetchall()
                result = [list(row) for row in result]
                field_names = self.getFunctionReturnFieldNames(function)
                result = [dict(zip(field_names, row)) for row in result]
                return result
            except:
                print("ERROR in functionCall: " + function + " does not return a single value or table")
                raise

    def procedureCall(self, procedure, named_args=None):
        '''
        calls the stored procedure given by "procedure" with arguments provided in dictionary "named_args"
        returns the return value of the procedure
        '''
        if named_args is not None:
            self.checkParameterNames(procedure, named_args.keys())
            # join named_args together into the appropriate query string:
            # eg., "exec Framework.sp param1 = 5, param2 = 'string'"
            parameter_names = self.getProcedureParameters(procedure)
            args = [named_args[parameter_name] for parameter_name in parameter_names]
            params_args = zip(parameter_names, args)
            for i, param_arg in enumerate(params_args):
                param = param_arg[0]
                arg = param_arg[1]

                if isinstance(arg, str) or isinstance(arg, unicode): 
                    arg = str(arg) # convert to plain string, in case unicode
                    start = 0
                    found = []
                    pat = re.compile("([^'])'([^'])")
                    while True:
                        m = pat.search(arg, start)
                        if m is None:
                            break
                        start = m.start() + 1
                        found.append(m.group(0))
                    for j in xrange(len(found)):
                        arg = re.sub(pat, "\\1''\\2", arg)

                    if re.search('\n\r', arg):
                        arg_split = arg.split('\n\r')
                        arg = '\\n\\r'.join(arg_split)
                        args[i] = quote(arg)
                    elif re.search('\n', arg):
                        arg_split = arg.split('\n')
                        arg = '\\n'.join(arg_split)
                        args[i] = quote(arg)
                    args[i] = quote(arg) # put strings in quotes
                    #args[i] = '%r' % arg
                elif isinstance(arg, datetime.datetime):
                    args[i] = arg.strftime('%m/%d/%Y %H:%M:%S')
                elif isinstance(arg, bool):
                    args[i] = str(int(arg))
                elif arg is not None:
                    args[i] = str(arg) # convert non-strings to strings
                else:
                    args[i] = 'NULL' # replace None with 'NULL'

                args[i] = param + ' = ' + args[i]

            #args = [arg for arg in args if arg != 'NULL']
            args = ', '.join(args)
        else:
            # if there are no arguments, form a simple call string
            args = ""
        
        procedure_string = procedure + " " + args

        self.cursor.execute("SET NOCOUNT ON")

        call = '\n'.join(["DECLARE @returnvalue int;", 
                          "EXEC @returnvalue = " + procedure_string + ";", 
                          "SELECT @returnvalue;",])
        try:
            self.cursor.execute(call)
        except:
            try: 
                self.execute('EXEC ' + procedure_string)
            except:
                type, value, tb = sys.exc_info()
                raise mssql_error(message = '\n'.join(["Error raised executing '" + procedure_string + "'.", str(type), str(value)]))

        result = self.cursor.fetchall()
        result = [list(row) for row in result]
        if len(result) > 0:
            if len(result[0]) == 1:
                result = [row[0] for row in result]
                #if len(result) == 1:
                #    result = result[0]
            else:
                result = [list(row) for row in result]

        self.cursor.execute("SET NOCOUNT OFF")

        return result

    def checkParameterNames(self, procedure, names):
        '''
        checks that the parameter names supplied in "names" match 
        the actual parameter names for stored procedure "procedure",
        raising an error if not the case.
        '''
        parameter_names = self.getProcedureParameters(procedure)
        try:
            assert set(names) == set(parameter_names)
        except AssertionError:
            print("Error: Parameter names provided do not match the stored procedure's parameter names.")
            print(procedure)
            print('Names supplied:')
            print(names)
            print('Actual procedure names:')
            print(parameter_names)
            raise

    def checkColumnNames(self, table, names):
        col_names = self.getColumnNames(table)
        try:
            assert set(names) == set(parameter_names)
        except AssertionError:
            print("Error: Column names provided do not match the table's column names.")
            print(table)
            print(names)
            print(col_names)
            raise

    def diffTable(self, table_name, ordered_by, rows, ignore=True):
        table_rows = self.getAllRows(table_name, order_by=ordered_by)
        if len(rows) == 0:
            return table_rows == rows
        else:
            rows_col_names = rows[0].keys()
            table_col_names = self.getColumnNames(table_name)
            differences = []
            if set(rows_col_names) != set(table_col_names):
                differences.append('column names')
            elif len(table_rows) != len(rows):
                differences.append('number of rows')
            if len(differences) > 0:
                return differences
            else:
                for i in xrange(len(rows)):
                    for col in table_col_names:
                        if rows[i][col] != table_rows[i][col]:
                            if rows[i][col] == "IGNORE":
                                next
                            else:
                                differences.append((i, col))
                return differences

    def coerceTSQL(self, value, tsql_type):
        if value == 'NULL': return None
        if value == "IGNORE": return "IGNORE"
        if tsql_type == 'bit': return bool(int(value))
        if tsql_type == 'datetime' or tsql_type == 'smalldatetime': return datetime.datetime.strptime(value, '%m/%d/%Y %H:%M')
        if tsql_type == 'float': return float(value)
        if tsql_type == 'int' or tsql_type == 'smallint': return int(value)
        if tsql_type == 'varchar': return str(value)

    def csvImportForExistingTable(self, csv_path, table_name):
        table_name_split = table_name.split('.')
        if len(table_name_split) > 1: table_name = table_name_split[1] # if table name includes schema prefix, remove it
        table_name = quote(table_name)
        self.cursor.execute('\n'.join(['SELECT DATA_TYPE',
                                       'FROM INFORMATION_SCHEMA.COLUMNS',
                                       'WHERE TABLE_NAME = ' + table_name,
                                       'ORDER BY ORDINAL_POSITION']))
        data_types = self.cursor.fetchall() 

        if len(data_types) == 0:
            sys.exit("Error: csvImportForExistingTable: Can't find " + table_name + '.')

        data_types = [str(row[0]) for row in data_types]

        try:
            f = open(csv_path)
        except:
            print "Can't open " + csv_path + ". No such file"
            raise
        reader = csv.reader(f, quotechar='|')
        header = reader.next()
        if len(data_types) != len(header):
            print table_name
            sys.exit("Error: DB.csvImportForExistingTable: Number of columns in CSV does not match the number of columns in the DB table")
        rows = []
        for line in reader:
            row = {}
            for i, colname in enumerate(header):
                row[colname] = self.coerceTSQL(line[i], data_types[i])
            rows = rows + [row]

        return rows

    def dictToXML(self, d):
        xml_str = '<ROOT>'

        for key, value in d.items():
            if value is None: 
                value = 'NULL'
            elif isinstance(value, str): 
                value = value.replace('&', '&amp;')
                value = value.replace('"', '&quot;')
                value = value.replace('<', '&lt;')
                value = value.replace('>', '&gt;')
                value = value.replace("'", '&apos;&apos;')
                value = "''" + value + "''"

            xml_str = xml_str + '<Element param="' + key + '" value="' + str(value) + '"></Element>'
        xml_str = xml_str + '</ROOT>'

        return xml_str

    def sp(self, procedure, **kwargs):
        final_args = None
        if len(kwargs) > 0:
            param_values = kwargs.values()
            param_keys = kwargs.keys()
            db_param_keys = ['@' + key for key in param_keys]
            procedure_param_names = self.getProcedureParameters(procedure)
            final_args = [(db_param, value) for db_param, value in zip(db_param_keys, param_values) if db_param in procedure_param_names]
            final_args = dict(final_args)
        
        return self.procedureCall(procedure, final_args)

    def ufn(self, function, **kwargs):
        final_args = None
        if len(kwargs) > 0:
            param_values = kwargs.values()
            param_keys = kwargs.keys()
            db_param_keys = ['@' + key for key in param_keys]
            function_param_names = self.getProcedureParameters(function)
            final_args = [(db_param, value) for db_param, value in
                          zip(db_param_keys, param_values) if db_param in
                          function_param_names]
            final_args = dict(final_args)
        
        return self.functionCall(function, final_args)

    def tableExists(table_name):
        tname_split = table_name.split('.')
        if len(tname_split) < 2:
            sys.exit('DB: tableExists: table schema not supplied')

        schema, table = tname_split[0], tname_split[1]
        self.cursor.execute('SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ' + quote(schema) + ' AND TABLE_NAME = ' + quote(table))
        
        return len(self.cursor.fetchall()) > 0

# misc support functions
def isNumber(s):
    # checks if a string is representing a number
    try:
        float(s)
        return True
    except ValueError:
        return False

def quote(s):
    return("'" + s + "'")

def truncToHour(dtime):
    return(datetime.strptime(datetime.strftime(dtime, "%Y-%m-%d %H"), "%Y-%m-%d %H"))

